import os
import scipy.io as sio
from torchvision.datasets import VisionDataset
from torchvision.datasets.folder import default_loader
from torchvision.datasets.utils import download_url
from torchvision.datasets.utils import extract_archive
from torchvision.transforms import transforms
from torch.utils.data.dataset import  Dataset


train_transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.RandomCrop((448, 448)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize((448, 448)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])


class CARDataSet(Dataset):
    def __init__(self, root, train=True):
        self.root = root
        self.loader = default_loader
        self.train = train

        loaded_mat = sio.loadmat(os.path.join(self.root, "cars_annos.mat"))
        loaded_mat = loaded_mat['annotations'][0]
        self.samples = []
        for item in loaded_mat:
            if self.train != bool(item[-1][0]):
                path = str(item[0][0])
                label = int(item[-2][0]) - 1
                self.samples.append((path, label))

    def __getitem__(self, index):
        path, target = self.samples[index]
        path = os.path.join(self.root, path)

        image = self.loader(path)
        if self.train:
            image = train_transform(image)
        else:
            image = test_transform(image)
        return image, target

    def __len__(self):
        return len(self.samples)


if __name__ == '__main__':
    from torch.utils.data.dataloader import DataLoader
    train_dataset = CARDataSet('../cars', train=True)
    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=False, num_workers=1)
    dataiter = iter(train_loader)
    images, labels = next(dataiter)
    print(images.shape)
    print(labels)

    test_dataset = CARDataSet('../cars', train=False)
    test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False, num_workers=1)
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    print(images.shape)
    print(labels)
